from typing import TypedDict
from langgraph.constants import END, START
from langgraph.graph import StateGraph


class InputState(TypedDict):
    """输入状态类型定义

    定义了图的输入数据结构
    """
    user_input: str


class OutputState(TypedDict):
    """输出状态类型定义

    定义了图的输出数据结构
    """
    graph_output: str


class OverallState(TypedDict):
    """整体状态类型定义

    定义了图中节点间传递的整体状态数据结构
    """
    foo: str
    user_input: str
    graph_output: str


class PrivateState(TypedDict):
    """私有状态类型定义

    定义了特定节点间使用的私有状态数据结构
    """
    bar: str


def node_1(state: InputState) -> OverallState:
    """处理节点1函数

    接收输入状态，将用户输入与固定字符串拼接后写入整体状态

    Args:
        state: 输入状态，包含用户输入的字符串

    Returns:
        更新后的整体状态，其中foo字段为用户输入加上">长沙市"
    """
    # Write to OverallState
    return {"foo": state["user_input"] + ">长沙市"}


def node_2(state: OverallState) -> PrivateState:
    """处理节点2函数

    读取整体状态中的foo字段，拼接固定字符串后写入私有状态

    Args:
        state: 整体状态，包含foo字段的数据

    Returns:
        私有状态，其中bar字段为foo字段值加上">岳麓区"
    """
    # Read from OverallState, write to PrivateState
    return {"bar": state["foo"] + ">岳麓区"}


def node_3(state: PrivateState) -> OutputState:
    """处理节点3函数

    读取私有状态中的bar字段，拼接固定字符串后写入输出状态

    Args:
        state: 私有状态，包含bar字段的数据

    Returns:
        输出状态，其中graph_output字段为bar字段值加上">欧富安科技园"
    """
    # Read from PrivateState, write to OutputState
    return {"graph_output": state["bar"] + ">欧富安科技园"}


# 构建状态图，指定整体状态、输入状态和输出状态类型
builder = StateGraph(
    OverallState,  # 整体状态类型
    input_schema=InputState,  # 输入状态类型
    output_schema=OutputState  # 输出状态类型
)

# 添加图节点，将函数与节点名称关联
builder.add_node("node_1", node_1)
builder.add_node("node_2", node_2)
builder.add_node("node_3", node_3)

# 添加节点间的连接边，定义执行流程
builder.add_edge(START, "node_1")
builder.add_edge("node_1", "node_2")
builder.add_edge("node_2", "node_3")
builder.add_edge("node_3", END)

# 编译构建好的图结构，生成可执行的图对象
graph = builder.compile()

# 调用图执行，传入初始输入状态
res = graph.invoke({"user_input": "湖南省"})
print(res)

from IPython.display import Image, display
# draw_mermaid⽅法可以打印出Graph的mermaid代码。
display(Image(graph.get_graph().draw_mermaid_png()))
